import os
import sys
from openpyxl import load_workbook
import xlsxwriter
from datetime import datetime
import traceback

def create_scatter_chart(workbook, worksheet, title, x_title, y_title, x_range, y_range, start_row=2, color=None):
    print(f"Creating scatter chart: {title}")
    chart = workbook.add_chart({'type': 'scatter'})
    chart.set_title({'name': title, 'name_font': {'size': 14, 'bold': True}})
    chart.set_x_axis({'name': x_title, 'num_format': 'hh:mm:ss', 'name_font': {'size': 12, 'bold': True}, 'num_font': {'size': 10}})
    chart.set_y_axis({'name': y_title, 'name_font': {'size': 12, 'bold': True}, 'num_font': {'size': 10}})
    chart.set_size({'width': 1200, 'height': 600})
    chart.set_legend({'position': 'none'})

    chart.add_series({
        'categories': [worksheet.name, start_row, x_range[0]-1, worksheet.dim_rowmax-1, x_range[0]-1],
        'values': [worksheet.name, start_row, y_range[0]-1, worksheet.dim_rowmax-1, y_range[0]-1],
        'marker': {'type': 'circle', 'size': 3},
        'line': {'color': color, 'width': 1.5}
    })

    return chart

def parse_time(time_str):
    try:
        return datetime.strptime(time_str.strip(), '%H:%M:%S')
    except ValueError:
        print(f"Warning: Unable to parse time: {time_str}")
        return None

def process_sheet_data(raw_wb, workbook, sheet_name):
    print(f"***Processing sheet: {sheet_name}")
    raw_sheet = raw_wb[sheet_name]
    worksheet = workbook.add_worksheet(f"{sheet_name}_Chart")

    # Define formats
    date_format = workbook.add_format({'num_format': 'hh:mm:ss', 'font_color': 'white'})
    hidden_format = workbook.add_format({'font_color': 'white'})
    sheet_name_format = workbook.add_format({
        'bold': True,
        'font_size': 32,
        'font_name': 'Arial',
        'align': 'right',
        'valign': 'bottom'
    })

    headers = next(raw_sheet.iter_rows(min_row=3, max_row=3, values_only=True))
    worksheet.write_row(0, 0, headers, hidden_format)

    row = 1
    for raw_row in raw_sheet.iter_rows(min_row=4, values_only=True):
        time = parse_time(raw_row[0])
        if time:
            worksheet.write_datetime(row, 0, time, date_format)
            for col, value in enumerate(raw_row[1:], start=1):
                worksheet.write(row, col, value, hidden_format)
            row += 1
        else:
            print(f"Skipping row with invalid time: {raw_row}")

    if row < 2:
        print(f"Warning: No valid data found in {sheet_name}")
        return

    print("Creating charts")
    charts = []
    colors = ['#4472C4', '#ED7D31', '#A5A5A5', '#FFC000', '#5B9BD5', '#70AD47', '#264478', '#9E480E']
    for i, header in enumerate(headers[1:], start=1):
        if f"{header}" != "None" and header != "CPU":
            color = colors[i % len(colors)]
            chart = create_scatter_chart(workbook, worksheet, f"{header} over Time", "Time", header, (1, 1), (i+1, i+1), start_row=1, color=color)
            chart.set_size({'width': 1000, 'height': 600})
            charts.append(chart)

    print("Adding charts to sheet\n")
    chart_width, chart_height = 17, 32
    charts_per_row = 3
    start_row = 0
    start_col = 0

    for i, chart in enumerate(charts):
        row = start_row + (i // charts_per_row) * chart_height
        col = start_col + (i % charts_per_row) * chart_width
        worksheet.insert_chart(row, col, chart)

    # Calculate the position for the sheet name
    last_chart_row = start_row + ((len(charts) - 1) // charts_per_row + 1) * chart_height
    last_chart_col = start_col + min(len(charts), charts_per_row) * chart_width

    # Add sheet name at the bottom right
    sheet_name_row = last_chart_row - 3
    sheet_name_col = last_chart_col - 3
    worksheet.write(sheet_name_row, sheet_name_col, sheet_name, sheet_name_format)

    # Set print area using xlsxwriter method (include sheet name)
    worksheet.print_area(start_row, start_col, sheet_name_row, last_chart_col)

    # Page setup for xlsxwriter
    worksheet.set_paper(9)  # A4 paper
    worksheet.fit_to_pages(1, 0)  # Fit to 1 page wide, height automatic
    worksheet.set_margins(left=0.3, right=0.3, top=0.3, bottom=0.3)
    worksheet.set_landscape()
    worksheet.hide_gridlines(2)  # Hide gridlines on screen and in print

    # Set zoom to 100% (xlsxwriter uses percentage directly)
    worksheet.set_zoom(100)

def main():
    if len(sys.argv) != 3:
        print("Usage: python generate_charts_from_raw.py <raw_data_xlsx> <output_dir>")
        sys.exit(1)

    raw_data_file = sys.argv[1]
    output_dir = sys.argv[2]

    if not os.path.isfile(raw_data_file):
        print(f"Error: {raw_data_file} is not a valid file.")
        sys.exit(1)

    os.makedirs(output_dir, exist_ok=True)

    try:
        print(f"Opening raw data file: {raw_data_file}")
        raw_wb = load_workbook(raw_data_file, read_only=True)

        output_file = os.path.join(output_dir, f'resource_charts_{datetime.now().strftime("%Y_%m_%d_%H_%M")}.xlsx')
        print(f"Creating new workbook: {output_file}")
        workbook = xlsxwriter.Workbook(output_file)

        process_sheet_data(raw_wb, workbook, 'CPU_ALL_RAW')
        process_sheet_data(raw_wb, workbook, 'MEM_RAW')

        # Process program-specific sheets
        for sheet_name in raw_wb.sheetnames:
            if sheet_name.endswith('_RAW') and sheet_name not in ['CPU_ALL_RAW', 'MEM_RAW']:
                process_sheet_data(raw_wb, workbook, sheet_name)

        print("Closing workbook")
        workbook.close()
        print(f"Chart file has been created: {output_file}")
    except Exception as e:
        print(f"An error occurred: {str(e)}")
        print("Traceback:")
        print(traceback.format_exc())
        sys.exit(1)

if __name__ == "__main__":
    main()
